import numpy as np
from DEAD.AutoDecoder.Evaluation.InternalBallistic import internal_ballistic_calculation, vectoring_internal_ballistic, generate_random, generate_random_log
from DEAD.AutoDecoder.Evaluation.Decoding import setup_dead_model, get_vector_from_index

def write_vector_to_npz(result_save_path,
                        global_index,
                        grain_index,
                        lv,
                        gv,
                        bvs,
                        avs1,
                        avs2):
    os.makedirs(f"{result_save_path}", exist_ok=True)
    info = [f"grain_index: {grain_index}", "lv: latent vector", "gv: geometry vector",
            "bvs: ballistic vectors", "avs1: attribute vectors 1", "avs2: attribute vectors 2"]
    np.savez_compressed(f"{result_save_path}/{global_index}.npz",
                        info=info,
                        lv=lv,
                        gv=gv,
                        bvs=bvs,
                        avs1=avs1,
                        avs2=avs2)


def random_internal_ballistic(current_item_index, w_array_norm, Ab_array_norm, Ap_array_norm):
    from DEAD.AutoDecoder.Config import cstar_min, cstar_max, rho_min, rho_max, r_ref_min, r_ref_max, n_min, n_max, R_min, R_max, J_min, J_max
    cstar_rho = generate_random(cstar_min*rho_min, cstar_max*rho_max)
    r_ref = generate_random(r_ref_min, r_ref_max)
    n = generate_random(n_min, n_max)
    R = generate_random_log(R_min, R_max)
    J = generate_random(J_min, J_max)  # 喉通比
    At = Ap_array_norm[0]*R*R*J
    t_array, p_array = internal_ballistic_calculation(
        w_array_norm=w_array_norm,
        Ab_array_norm=Ab_array_norm,
        cstar_rho=cstar_rho,
        r_ref=r_ref,
        n=n,
        R=R,
        At=At)
    if t_array is not None and p_array is not None:
        t_action, t_vector, p_vector = vectoring_internal_ballistic(
            t_array, p_array)
        bv = np.stack([t_vector, p_vector*1e-6], axis=1)
        av1= np.array([(r_ref-r_ref_min)/(r_ref_max-r_ref_min),  # 推进剂燃烧（燃速）特性
                       (R-R_min)/(R_max-R_min)  # 装药半径
                       ])*0.1
        av2 = np.array([(cstar_rho-cstar_min*rho_min)/(cstar_max*rho_max-cstar_min*rho_min),  # 推进剂能量特性
                        (n-n_min)/(n_max-n_min),  # 推进剂压力特性
                       At*100.0])  # 喉部面积
        next_item_index = current_item_index+1
        return next_item_index, bv, av1, av2
    else:
        next_item_index = current_item_index
        return next_item_index, None, None, None


def task(dead, decode_result_save_path, PSFD_save_path, npz_files, grain_index, item_number_per_grain):
    lv, gv = get_vector_from_index(dead, grain_index)
    lv = lv[0].detach().cpu().numpy()
    gv = gv[0].detach().cpu().numpy()
    npz_file = [file for file in npz_files if file.startswith(
        f'{grain_index}#')][0]
    npz_file = np.load(f"{decode_result_save_path}/{npz_file}")
    w_array_norm = npz_file['burned_web_array']
    Ab_array_norm = npz_file['burning_surface_area_array']
    Ap_array_norm = npz_file['exit_port_area_array']

    current_item_index = 0
    bvs=[]
    avs1=[]
    avs2=[]
    i=0
    while (current_item_index < item_number_per_grain) and (i<1e6):#防止超时
        current_item_index, bv, av1 ,av2 = random_internal_ballistic(
            current_item_index, w_array_norm, Ab_array_norm, Ap_array_norm)
        i+=1
        if av2 is not None:
            bvs.append(bv)
            avs1.append(av1)
            avs2.append(av2)

    if not avs2: #列表是空的
        print("grain_index:{grain_index} 没有合适的燃速产生合理的内弹道曲线")
    else: #列表不是空的
        global global_index
        bvs = np.array(bvs)
        avs1 = np.array(avs1)
        avs2 = np.array(avs2)
        write_vector_to_npz(PSFD_save_path,
                            global_index,
                            grain_index,
                            lv,
                            gv,
                            bvs,
                            avs1,
                            avs2)
        global_index += 1


if __name__ == '__main__':
    from DEAD.AutoDecoder.Config import *
    item_number_per_grain = 100

    dead = setup_dead_model(hidden_size, depth, data_num,
                            latent_size, model_load_path)
    dead.cuda()

    # Get all files in the directory
    files = os.listdir(f"{decode_result_save_path}")
    npz_files = [file for file in files if file.endswith('.npz')]
    global_index=0
    for grain_index in range(0, data_num):
        task(dead, decode_result_save_path, PSFD_save_path,
             npz_files, grain_index, item_number_per_grain)
        print(f"finish {grain_index+1}/{data_num}")
